from fund.strategy.strategy_util import process_dates, empty_result
from util.csv_util import csv_2_df
from util.date_util import *

def strategy_based_on_moving_average(df, higher_money=0, lower_money=200, ma_range=30, start_date=None, end_date=None):
    """
    均值策略（增强版）：记录每次买入的日期和价格，便于后续可视化
    核心逻辑：
    1. 计算前 days 天的净值均值
    2. 当天价格 >= 均线时买入 x 元，否则买入 y 元
    3. 记录每次买入的日期、价格、金额和份额

    :param df: 基金净值数据（需包含日期索引和'单位净值'列）
    :param higher_money: 价格高于均线时买入金额
    :param lower_money: 价格低于均线时买入金额
    :param ma_range: 均线计算窗口
    :param start_date: 策略开始日期
    :param end_date: 策略结束日期
    :return: 包含收益指标和买入记录的字典
    """
    # --- 初始化 ---
    start_date, end_date = process_dates(df, start_date, end_date)
    total_cost = 0
    total_units = 0
    attempts_times = 0

    # 新增：记录每次买入操作
    buy_records = []  # 存储每次买入的信息

    # 获取日期索引位置
    start_index = df.index.searchsorted(start_date)
    end_index = df.index.searchsorted(end_date)

    # 确保足够计算均线
    start_index = max(start_index, ma_range)  # 必须满足 start_index >= days

    if start_index >= end_index:
        return empty_result(start_date, end_date)

    # --- 核心逻辑 ---
    for current_pos in range(start_index, end_index + 1):
        current_net_value = df['单位净值'].iloc[current_pos]
        current_date = df.index[current_pos]  # 获取当前日期

        # 计算均线
        last_days_mean = df['单位净值'].iloc[current_pos - ma_range: current_pos].mean()

        # 决策买入金额
        if current_net_value >= last_days_mean:
            daily_cost = higher_money
        else:
            attempts_times += 1
            daily_cost = lower_money

        # 计算买入份额
        daily_units = daily_cost / current_net_value
        total_units += daily_units
        total_cost += daily_cost

        # 新增：记录买入操作
        if daily_cost > 0:  # 只记录实际买入的操作
            buy_records.append({
                'date': current_date,
                'price': current_net_value,
                'cost': daily_cost,
                'units': daily_units,
            })

    # --- 收益计算 ---
    if end_index > start_index:
        last_net_value = df['单位净值'].iloc[end_index]
        final_asset = total_units * last_net_value
    else:
        final_asset = 0

    total_profit = final_asset - total_cost
    total_return_rate = total_profit / total_cost if total_cost != 0 else 0
    investment_days = (end_date - start_date).days

    # 年化收益率
    if total_cost != 0 and investment_days > 0:
        annual_return_rate = (final_asset / total_cost) ** (365 / investment_days) - 1
    else:
        annual_return_rate = 0

    return {
        "开始时间": start_date.strftime('%Y-%m-%d'),
        "结束时间": end_date.strftime('%Y-%m-%d'),
        "总投入": total_cost,
        "最终资产": final_asset,
        "总利润": total_profit,
        "总收益率": total_return_rate,
        "年化收益率": annual_return_rate,
        "买入次数": attempts_times,
        "买入记录": buy_records
    }

# ---------------------- 示例：如何使用这个函数 ----------------------
if __name__ == "__main__":
    # 1. 读取基金数据
    # fund_code = '022364'
    # fund_code = '580006'
    fund_code = '007540'
    df = csv_2_df(f"fund/history_data/{fund_code}-daily.csv")

    res = strategy_based_on_moving_average(df, start_date=before_2y_yyyymmdd())
    # res = strategy_based_on_moving_average(df, start_date=before_1y_yyyymmdd())
    print(res['买入记录'][-1])
    print(len(res['买入记录']))